#include "kernel_operator.h"
#include "utils.h"
using namespace AscendC;

extern "C" __global__ __aicore__ void gemm_a310p(
  const int m, const int n, const int k,
  const uint8_t * __restrict__ A, const int lda,
  const uint8_t * __restrict__ B, const int ldb,
  int32_t * __restrict__ C, const int ldc,
  const int bm, const int bn, const int bk, const int am, const int an)
{
  const int queue_buf_sz = 2;
  const int aicore_id = GetBlockIdx();
  TPipe pipe;
  TQue<QuePosition::VECIN, queue_buf_sz> queueA;
  TQue<QuePosition::VECIN, queue_buf_sz> queueB;
  TQue<QuePosition::VECOUT, queue_buf_sz> queueC;
  pipe.InitBuffer(queueA, queue_buf_sz, src_tw*src_th*sizeof(char));
  pipe.InitBuffer(queueB, queue_buf_sz, ((src_tw*src_th)/4)*sizeof(char));
  pipe.InitBuffer(queueC, queue_buf_sz, tw*th*sizeof(float));

  for(int tx = 0; tx < bw; tx+=tw){
    for(int ty = 0; ty < bh; ty+=th){
      Rect clip_tile_st_roi, clip_tile_ld_roi;
      Rect tile_st_roi = {b_idx*bw+tx, b_idy*bh+ty, tw, th};
      inter(tile_st_roi, dstRoi, clip_tile_st_roi);
      tile_st_roi = clip_tile_st_roi;
      Rect tile_ld_roi = {
        floor((tile_st_roi.x-dstRoi.x+0.5f)*scale_x - 0.5f), 
        floor((tile_st_roi.y-dstRoi.y+0.5f)*scale_y - 0.5f),
        floor(tile_st_roi.width*scale_x+1.0f),
        floor(tile_st_roi.height*scale_y+1.0f)
      };
      inter(tile_ld_roi, srcRoi, clip_tile_ld_roi);
      tile_ld_roi = clip_tile_ld_roi;
      if(tile_st_roi.width == 0 || 
        tile_st_roi.height == 0 ||
        tile_ld_roi.width == 0 ||
        tile_ld_roi.height == 0){
        continue;
      }
      //load src global
      LocalTensor<char> srcYLocal = inQueueSrcY.AllocTensor<char>(), srcUvLocal = inQueueSrcUv.AllocTensor<char>();
      for(int ty_tid = tile_ld_roi.y; ty_tid < tile_ld_roi.height; ++ty_tid){ //load srcY
        GlobalTensor<char> src_gts;
        src_gts.SetGlobalBuffer(srcy+ty_tid*src_sw+tile_ld_roi.x, tile_ld_roi.width);
        DataCopy(srcYLocal, src_gts, tile_ld_roi.width);
      }
      for(int th_tid = ty; th_tid < (src_th/2); ++th_tid){ //load srcUv
        int h_id = 
        GlobalTensor<char> src_gts;
        src_gts.SetGlobalBuffer(srcuv + , block_sz);
      }
      //math
      //store dst global
    }
  }

  src_gts.SetGlobalBuffer(x + block_sz * GetBlockIdx(), block_sz);
  dst_gts.SetGlobalBuffer(y + block_sz * GetBlockIdx(), block_sz);
  // printf("%f\n", xGm.GetValue(0));
  pipe.InitBuffer(inQueueX, queue_buf_sz, tile_sz * sizeof(float));
  pipe.InitBuffer(inQueueY, queue_buf_sz, tile_sz * sizeof(float));
  pipe.InitBuffer(outQueueZ, queue_buf_sz, tile_sz * sizeof(float));
  int32_t block_tsz =  block_sz/tile_sz;
  for (int32_t i = 0; i < block_tsz; i++) {
    LocalTensor<float> xLocal = inQueueX.AllocTensor<float>(), yLocal = inQueueY.AllocTensor<float>();
    DataCopy(xLocal, xGm[i*tile_sz], tile_sz);
    DataCopy(yLocal, yGm[i*tile_sz], tile_sz);
    inQueueX.EnQue(xLocal);
    inQueueY.EnQue(yLocal);
    xLocal = inQueueX.DeQue<float>();
    yLocal = inQueueY.DeQue<float>();
    LocalTensor<float> zLocal = outQueueZ.AllocTensor<float>();
    AscendC::Add(zLocal, xLocal, yLocal, tile_sz);
    outQueueZ.EnQue<float>(zLocal);
    inQueueX.FreeTensor(xLocal);
    inQueueY.FreeTensor(yLocal);
    zLocal = outQueueZ.DeQue<float>();
    DataCopy(zGm[i * tile_sz], zLocal, tile_sz);
    outQueueZ.FreeTensor(zLocal);
  }
  // printf("z %f\n", z[0]);
}

/**
 * Weight: A (float32)
 * Input Activation: B (float32)
 * Output Activation: C (float32)
*/
int gemm(
  const int m, const int n, const int k,
  float *A, const int lda, 
  float *B, const int ldb, 
  float *C, const int ldc,
  void *stream = nullptr)
{
  const int bm = 128, bn = 128;
  const int am = DIV_UP(m, bw), an = DIV_UP(n, bn);
  const int queue_buf_sz = 2;
  gemm_a310p<<<aw*ah, nullptr, stream>>>(
    m, n, k, A, lda, B, ldb, C, ldc,
    bm, bn, am, an);
  return 0;
}
